import pandas as pd
from sklearn.model_selection import train_test_split
import os

def main():
    # 读取数据
    df = pd.read_csv("./dataset/small_data.csv")

    # 重命名列名以确保一致性
    df = df.rename(columns={"text": "text", "label": "label"})

    # 去除缺失值
    df = df.dropna(subset=['text', 'label'])

    # 确保 label 列是整数类型
    df['label'] = df['label'].astype(int)

    # 打乱数据
    df = df.sample(frac=1, random_state=42)

    # 划分数据集
    train, val_test = train_test_split(df, test_size=0.2, random_state=42)
    val, test = train_test_split(val_test, test_size=0.5, random_state=42)

    # 创建 dataset 文件夹（如果不存在）
    os.makedirs("dataset", exist_ok=True)

    # 保存到 CSV 文件
    train.to_csv("dataset/train.csv", index=False)
    val.to_csv("dataset/val.csv", index=False)
    test.to_csv("dataset/test.csv", index=False)

    print("数据集划分完成！样本量：")
    print(f"训练集：{len(train)}")
    print(f"验证集：{len(val)}")
    print(f"测试集：{len(test)}")

if __name__ == "__main__":
    main()